package middleware

import (
	"fastgin/utils/res"
	"time"

	"github.com/gin-gonic/gin"
)

func LimitMiddleware(limit int) gin.HandlerFunc {
	return NewLimiter(limit, 1*time.Second).Middleware
}
func NewLimiter(limit int, duration time.Duration) *Limiter {
	return &Limiter{
		limit:     limit,
		duration:  duration,
		timestamp: make(map[string][]int64),
	}
}

// 限流器
type Limiter struct {
	limit     int
	duration  time.Duration
	timestamp map[string][]int64
}

func (l *Limiter) Middleware(c *gin.Context) {
	//获取客户端ip地址
	ip := c.ClientIP()
	//检查请求时间戳切片是否存在
	if _, ok := l.timestamp[ip]; !ok {
		l.timestamp[ip] = make([]int64, 0)
	}
	//当前时间戳
	now := time.Now().Unix()
	//移除过期时间戳
	for i := 0; i < len(l.timestamp[ip]); i++ {
		if l.timestamp[ip][i] < now-int64(l.duration.Seconds()) {
			l.timestamp[ip] = append(l.timestamp[ip][:i], l.timestamp[ip][i+1:]...)
			i--
		}
	}
	//检查请求数量是否超过限制
	if len(l.timestamp[ip]) > l.limit {
		res.FailWithMsg("To Many Request", c)
		c.Abort()
		return
	}
	//添加当前请求时间戳到切片
	l.timestamp[ip] = append(l.timestamp[ip], now)
	//继续处理请求
	c.Next()
}
